package sshex

import (
	"context"
	"encoding/hex"
	"net"

	"golang.org/x/crypto/ssh"
)

var (
	ContextKeyUser          = "user"
	ContextKeySessionID     = "session-id"
	ContextKeyRole          = "role"
	ContextKeyClientVersion = "client-version"
	ContextKeyServerVersion = "server-version"
	ContextKeyLocalAddr     = "local-address"
	ContextKeyRemoteAddr    = "remote-address"
	ContextKeyPermissions   = "permission"
	ContextKeyPublicKey     = "public-key"
)

type ConnContext struct {
	context.Context
	cancel context.CancelFunc
}

func NewConnContext(ctx context.Context, con net.Conn) (*ConnContext, context.CancelFunc) {
	innerCtx, cancel := context.WithCancel(ctx)
	c := &ConnContext{
		Context: innerCtx,
		cancel:  cancel,
	}
	perms := &ssh.Permissions{}
	c.SetValue(ContextKeyPermissions, perms)
	return c, cancel
}

func (ctx *ConnContext) SetValue(key, value interface{}) {
	ctx.Context = context.WithValue(ctx.Context, key, value)
}

func (ctx *ConnContext) ApplyConnMetadata(conn ssh.ConnMetadata) {
	if ctx.Value(ContextKeySessionID) != nil {
		return
	}
	ctx.SetValue(ContextKeySessionID, hex.EncodeToString(conn.SessionID()))
	ctx.SetValue(ContextKeyClientVersion, string(conn.ClientVersion()))
	ctx.SetValue(ContextKeyServerVersion, string(conn.ServerVersion()))
	ctx.SetValue(ContextKeyUser, conn.User())
	ctx.SetValue(ContextKeyLocalAddr, conn.LocalAddr())
	ctx.SetValue(ContextKeyRemoteAddr, conn.RemoteAddr())
}

func (ctx *ConnContext) User() string {
	return ctx.Value(ContextKeyUser).(string)
}
func (ctx *ConnContext) SessionID() string {
	return ctx.Value(ContextKeySessionID).(string)
}

func (ctx *ConnContext) ClientVersion() string {
	return ctx.Value(ContextKeyClientVersion).(string)
}

func (ctx *ConnContext) ServerVersion() string {
	return ctx.Value(ContextKeyServerVersion).(string)
}

func (ctx *ConnContext) RemoteAddr() net.Addr {
	return ctx.Value(ContextKeyRemoteAddr).(net.Addr)
}

func (ctx *ConnContext) LocalAddr() net.Addr {
	return ctx.Value(ContextKeyLocalAddr).(net.Addr)
}

func (ctx *ConnContext) Permissions() *ssh.Permissions {
	return ctx.Value(ContextKeyPermissions).(*ssh.Permissions)
}
